"""
训练MHD模型示例
"""

import os
import sys
import argparse

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from em_generation.config import get_config
from em_generation.train import main as train_main


def main():
    parser = argparse.ArgumentParser(description='训练MHD模型示例')
    parser.add_argument('--data_path', type=str, 
                        default='./data/synthetic/mhd_data.h5', 
                        help='数据文件路径')
    parser.add_argument('--output_dir', type=str, default='./outputs/mhd', 
                        help='输出目录')
    parser.add_argument('--epochs', type=int, default=50, 
                        help='训练轮数')
    parser.add_argument('--batch_size', type=int, default=4, 
                        help='批次大小')
    parser.add_argument('--lr', type=float, default=1e-3, 
                        help='学习率')
    parser.add_argument('--modes', type=int, default=12, 
                        help='傅里叶模式数')
    parser.add_argument('--width', type=int, default=32, 
                        help='网络宽度')
    parser.add_argument('--device', type=str, default='auto', 
                        help='计算设备')
    parser.add_argument('--seed', type=int, default=42, 
                        help='随机种子')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 构建训练参数
    train_args = [
        '--data_path', args.data_path,
        '--output_dir', args.output_dir,
        '--model_type', 'mhd',
        '--epochs', str(args.epochs),
        '--batch_size', str(args.batch_size),
        '--lr', str(args.lr),
        '--modes', str(args.modes),
        '--width', str(args.width),
        '--device', args.device,
        '--seed', str(args.seed)
    ]
    
    # 模拟命令行参数
    sys.argv = ['train_mhd_example.py'] + train_args
    
    print("开始训练MHD模型...")
    print(f"数据路径: {args.data_path}")
    print(f"输出目录: {args.output_dir}")
    print(f"训练轮数: {args.epochs}")
    print(f"批次大小: {args.batch_size}")
    print(f"学习率: {args.lr}")
    
    # 调用训练主函数
    train_main()


if __name__ == "__main__":
    main()